import os
import sys
import torch
import numpy as np
import cv2
import glob
import time
import shutil

# 添加项目根目录到系统路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

# 设置临时目录
def get_temp_dir():
    """创建临时目录用于存储中间结果"""
    # 使用tempfile模块创建临时目录，避免中文路径问题
    import tempfile
    temp_dir = tempfile.mkdtemp(prefix="udis2_temp_")
    return temp_dir

# 完整的UDIS2拼接流程
def stitch_images_udis2(img1, img2, gpu='0'):
    """使用UDIS2算法拼接两张图像"""
    print("开始UDIS2图像拼接流程...")
    
    # 记录开始时间
    start_time = time.time()
    
    # 创建临时目录
    temp_dir = get_temp_dir()
    print(f"使用临时目录: {temp_dir}")
    
    # 保存输入图像
    img1_path = os.path.join(temp_dir, 'img1.jpg')
    img2_path = os.path.join(temp_dir, 'img2.jpg')
    
    # 使用PIL库保存图像，避免中文路径问题
    try:
        from PIL import Image
        Image.fromarray(cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)).save(img1_path)
        Image.fromarray(cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)).save(img2_path)
        print(f"图像已保存到: {img1_path} 和 {img2_path}")
    except Exception as e:
        print(f"使用PIL保存图像失败: {str(e)}，尝试使用OpenCV")
        cv2.imwrite(img1_path, img1)
        cv2.imwrite(img2_path, img2)
    
    # 设置UDIS2 Warp阶段的路径
    warp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'UDIS2', 'Warp', 'Codes')
    warp_script = os.path.join(warp_dir, 'test_other.py')
    
    # 保存当前工作目录
    original_dir = os.getcwd()
    
    try:
        # 执行UDIS2 Warp阶段
        print("执行UDIS2 Warp阶段...")
        os.chdir(warp_dir)  # 切换到Warp目录
        cmd = f'python {warp_script} --gpu {gpu} --path {temp_dir}/ --img1_name img1.jpg --img2_name img2.jpg'
        print(f"执行命令: {cmd}")
        os.system(cmd)
        
        # 检查是否生成了必要的文件
        warp_files = ['warp1.jpg', 'warp2.jpg', 'mask1.jpg', 'mask2.jpg']
        for file in warp_files:
            file_path = os.path.join(temp_dir, file)
            if not os.path.exists(file_path):
                print(f"UDIS2 Warp阶段失败，未生成{file}，路径: {file_path}")
                os.chdir(original_dir)  # 恢复原始工作目录
                return None, None, {}
            else:
                print(f"成功生成文件: {file_path}")
        
        # 设置UDIS2 Composition阶段的路径
        comp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'UDIS2', 'Composition', 'Codes')
        comp_script = os.path.join(comp_dir, 'test_other.py')
        
        # 执行UDIS2 Composition阶段
        print("执行UDIS2 Composition阶段...")
        os.chdir(comp_dir)  # 切换到Composition目录
        cmd = f'python {comp_script} --gpu {gpu} --path {temp_dir}/'
        print(f"执行命令: {cmd}")
        os.system(cmd)
    finally:
        # 恢复原始工作目录
        os.chdir(original_dir)
    
    # 检查是否生成了最终结果
    composition_path = os.path.join(temp_dir, 'composition.jpg')
    if not os.path.exists(composition_path):
        print("UDIS2 Composition阶段失败，未生成最终拼接结果")
        return None, None, {}
    
    # 读取最终结果
    result_image = cv2.imread(composition_path)
    
    # 收集处理步骤
    processing_steps = [
        {
            "name": "输入图像1",
            "path": img1_path
        },
        {
            "name": "输入图像2",
            "path": img2_path
        },
        {
            "name": "变形图像1",
            "path": os.path.join(temp_dir, "warp1.jpg")
        },
        {
            "name": "变形图像2",
            "path": os.path.join(temp_dir, "warp2.jpg")
        },
        {
            "name": "掩码1",
            "path": os.path.join(temp_dir, "mask1.jpg")
        },
        {
            "name": "掩码2",
            "path": os.path.join(temp_dir, "mask2.jpg")
        },
        {
            "name": "学习掩码1",
            "path": os.path.join(temp_dir, "learn_mask1.jpg")
        },
        {
            "name": "学习掩码2",
            "path": os.path.join(temp_dir, "learn_mask2.jpg")
        },
        {
            "name": "彩色融合",
            "path": os.path.join(temp_dir, "composition_color.jpg")
        },
        {
            "name": "最终结果",
            "path": composition_path
        }
    ]
    
    # 计算执行时间
    execution_time = time.time() - start_time
    
    # 性能数据
    performance_data = {
        "execution_time": execution_time,
        "algorithm": "UDIS2"
    }
    
    return result_image, processing_steps, performance_data

# 清理临时文件
def cleanup_temp_files():
    """清理临时文件"""
    import tempfile
    import glob
    
    # 清理所有以udis2_temp_开头的临时目录
    for temp_dir in glob.glob(os.path.join(tempfile.gettempdir(), "udis2_temp_*")):
        if os.path.exists(temp_dir):
            try:
                print(f"清理临时目录: {temp_dir}")
                shutil.rmtree(temp_dir)
            except Exception as e:
                print(f"清理临时目录失败: {str(e)}")
    
    # 确保返回到项目根目录
    os.chdir(project_root)